import os
import codecs
import Levenshtein

base_dir = "/data/ppocr/trainData/test"


def rec_metrics(preds, labels):
    correct_num = 0
    all_num = 0
    norm_edit_dis = 0.0
    for pred, target in zip(preds, labels):
        pred = pred.replace(" ", "")
        target = target.replace(" ", "")
        norm_edit_dis += Levenshtein.distance(pred, target) / max(
            len(pred), len(target), 1)
        if pred == target:
            correct_num += 1
        all_num += 1
    return {
        'acc': correct_num / all_num,
        'norm_edit_dis': 1 - norm_edit_dis / all_num
    }

# {'acc': 0.7500893388921978, 'norm_edit_dis': 0.9225230862977484}


if __name__ == "__main__":
    truth_dict = {}
    pred_dict = dict()
    with codecs.open(os.path.join(base_dir, "truth_ocr.txt"), "r", "utf8") as f1:
        for line in f1.readlines():
            line = line.strip()
            if not line:
                continue
            split_line = line.split("\t")
            truth_dict[split_line[0]] = split_line[1]

    with codecs.open(os.path.join(base_dir, "paddle_ocr.txt"), "r", "utf8") as f2:
        for line in f2.readlines():
            line = line.strip()
            if not line:
                continue
            split_line = line.strip().split("\t")
            if len(split_line) == 1:
                continue
            pred_dict[split_line[0]] = split_line[1]

    target_list = []
    pred_list = []
    for key, val in truth_dict.items():
        target_list.append(val)
        if key in pred_dict.keys():
            pred_list.append(pred_dict[key])
        else:
            pred_list.append("")

    print(rec_metrics(target_list, pred_list))


# zhangbo  {'acc': 0.7161405598570578, 'norm_edit_dis': 0.9310714570819941}
# chinese lite {'acc': 0.7499702203692674, 'norm_edit_dis': 0.9224039677748181}
